"""
Filename: ss_gwaste.jl
Paper: Unemployment and the Distribution of Liquidity
Authors: Zach Bethune and Guillaume Rocheteau
Contact: bethune@rice.edu
Last Modified: 8/22/2023
Purpose: functions to compute a steady-state equilibrium when money creation is spent on wasteful government purchases.
"""

###############################################################################
# Market clearing conditions and function to compute equilibrium W'(a)
###############################################################################
#market clearing for (Rf,py) in steady state (Rm exog.)
function mkt_clear_ss!(prices; mm::model_outcomes,parms::parmsmodel,md::moments_data,W_iter_tol=1e-5,Zmd::Real)
    """
    purpose: Solve for steady state equilibrium prices
    inputs: initial guess of prices -> prices
    outputs: vector of market clearing with dimension 2 (markets for early consumption and financial wealth) *note: market for liquid wealth is imposed to clear in ss with assumption Rm = 1/(1+pi) exogenous
    """
    py = copy(prices[1]);
    Rf = copy(prices[2]);

    Zmd0=copy(Zmd); Zmd1=copy(Zmd);
    error=1.0; count_iter=0;

    while error>1e-3 && count_iter<10
        count_iter+=1;
        
        #taxes
        mm.tau=zeros((Na,2,Nz));
        taulumpsum = ((1.0/Rf)-1.0)*mm.Ag;
        mm.tau .= taulumpsum;

        #firm problem
        mm.Ys = κ_prime_inv(py;parms);
        mm.frev = mm.zgrid.*(1.0 + py*mm.Ys - κ(mm.Ys;parms));
        mm.wages_bar = zeros((2,Nz));
        mm.wages_bar[1,:] = md.LABOR_SHARE.*mm.frev;
        mm.wages_bar[2,:] = md.REPLACE_RATE.*mm.wages_bar[1,:]
            
        mm.wages = zeros((Na,2,Nz));
        for j=1:Nz
            for i=1:2
                mm.wages[:,i,j] .= mm.wages_bar[i,j].+mm.tau[:,i,j];
            end
        end

        profits=mm.frev.-mm.wages_bar[1,:]
        mm.Js = (profits)./(1.0-((1.0-parms.DELTA)/Rf))
        
        if maximum((mm.zgrid.*parms.K.*Rf)./mm.Js) < 1
            mm.theta = q_inv(((mm.zgrid.*parms.K*Rf)./mm.Js)[1];parms)
        else
            mm.theta = 0.0001
        end
        mm.emp = λ(mm.theta;parms)/(parms.DELTA+λ(mm.theta;parms));
        
        #household problem
        Wa_temp=deepcopy(mm.Wa);
        Wa_solve!(mm;tol=W_iter_tol,py,Rf,Wa_p=Wa_temp,mm.agrid,parms);
    
        transition_g!(mm;py,parms);   
        mm.g=zeros((Na,2,Nz));
        (lam,ghat)=powm!(mm.T,rand(Na*2*Nz),tol=1e-10,maxiter=1000000);
        mm.g = reshape(ghat,(Na,2,Nz))
        for j=1:Nz
            mm.g[:,:,j] = mm.g[:,:,j]*(mm.gz[j]/sum(mm.g[:,:,j]))
        end
        mm.ga = sum(mm.g,dims=3)[:,:,1]
        
        #compute market clearing condition
        #early consumption goods 
        mm.Yd=0.0
        for l=1:Nz
            for k=1:2
                for kk=1:2
                    ystar_int = LinearInterpolation(mm.agrid, mm.ystar[:,kk,l], extrapolation_bc=Line());
                    for j=1:Na
                        mm.Yd += mm.PI[k,kk]*parms.ALPHA*parms.ALPHAmf*ymf(ystar_int,mm.a_p[j,k,l],mm.am_p[j,k,l],py)[1]*mm.g[j,k,l]
                        mm.Yd += mm.PI[k,kk]*parms.ALPHA*(1.0-parms.ALPHAmf)*ym(ystar_int,mm.a_p[j,k,l],mm.am_p[j,k,l],py)[1]*mm.g[j,k,l]
                    end
                end
            end
        end
    
        #financial wealth 
        mm.Jd = sum(mm.g.*(mm.a_p.-mm.am_p))

        #liquid wealth
        Zmd1 = sum(mm.g.*mm.am_p);
        error = abs(Zmd0.-Zmd1);
        Zmd0=copy(Zmd1);
        mm.Zmd = copy(Zmd1);
    end

    z = zeros(2)
    z[1] = mm.emp*sum(mm.zgrid.*mm.gz)*mm.Ys - mm.Yd
    z[2] =  mm.Jd - mm.emp*sum(mm.gz.*mm.Js) - mm.Ag

    #record prices
    mm.py = copy(py)
    mm.Rf = copy(Rf)

    return z
end

#iterate Euler to solve for steady state equilibrium Wa, etc., given prices
function Wa_solve!(mm::model_outcomes;tol,py,Rf,Wa_p,agrid,parms::parmsmodel)
    """
    purpose: Solve for steady state W'(a), policies am', a', c, ystar 
    inputs: py, Rf, Rm, tau, wages, theta (all in steady state)
    outputs: Wa, am_p, a_p, c, ystar
    """
    mm.Wa=zeros((Na,2,Nz)); mm.am_p=zeros((Na,2,Nz)); mm.a_p=zeros((Na,2,Nz)); mm.c=zeros((Na,2,Nz)); mm.ystar=zeros((Na,2,Nz)); 
    error=1.0; iter=0; 
    while error>tol && iter<Wa_iter_max
        iter=iter+1;
        mm.Wa[:,:,:], mm.am_p[:,:,:], mm.a_p[:,:,:], mm.c[:,:,:], mm.ystar[:,:,:] = Wa_iter(;py,Rf,Wa_p,agrid,parms,mm)
        error = maximum(abs.(mm.Wa.-Wa_p))
        Wa_p = deepcopy(mm.Wa);
    end
end

#iterate Euler equations using endog. grid method
function Wa_iter(; py,Rf,Wa_p,agrid,parms::parmsmodel,mm::model_outcomes)
    """
    purpose: Given W'_{t+1}(a), update W'_t(a) and policies
    inputs: wages_t, py_{t+1}, Rf_t, Rm_t, tau_t, theta_{t+1}
    outputs: Wa_t and policies: am_{t+1}, a_{t+1}, c_t, ystar_{t+1}
    """
    #update post transfer wages
    mm.wages = zeros((Na,2,Nz));
    for j=1:Nz
        for i=1:2
            mm.wages[:,i,j] .= mm.wages_bar[i,j].+mm.tau[:,i,j];
        end
    end

    #employment transition matrix
    mm.PI = [(1.0-parms.DELTA) parms.DELTA; λ(mm.theta;parms) 1.0-λ(mm.theta;parms)]

    #Compute ystar(e,a) given py and Wa_p
    mm.ystar = ystar_update(py,Wa_p,agrid;parms); #root_finding

    #Update Vaj_p 
    Vam_p, Vaf_p = Vaj_update(py,mm.ystar,Wa_p,agrid;parms);

    #################################################
    #Decision rules am_p, a_p, c when unconstrained
    #################################################
    #define Xj = Rj*beta*Vaj_p
    Xam = zeros((Na,Na,2,Nz)); Xaf = zeros((Na,Na,2,Nz));
    
    for i=1:Nz
        Xam[:,:,:,i] = mm.Rm.*parms.BETA.*reshape((mm.PI*reshape(Vam_p[:,:,:,i],Na*Na,2)')',(Na,Na,2));
        Xaf[:,:,:,i] = Rf.*parms.BETA.*reshape((mm.PI*reshape(Vaf_p[:,:,:,i],Na*Na,2)')',(Na,Na,2));
    end

    #invert to get am_prime
    am_p_temp, c_temp = policy_update_step1(Xam,Xaf,agrid;parms); #root_finding

    #solve for a_p, am_p, c in unconstrained eq
    am_p_u, a_p_u, c_u = policy_update_step2(Rf,am_p_temp,c_temp,agrid;parms,mm); #no root_finding
    
    #################################################
    #Decision rules am_p, a_p, c when constrained
    #################################################
    am_p_c, a_p_c, c_c = policy_update_constrained(Xam,agrid;parms);

    #################################################
    #Optimal decision rules am_p, a_p, c
    #################################################
    mm.am_p = am_p_u.*(am_p_u.<am_p_c) .+ am_p_c.*(am_p_u.>=am_p_c);
    mm.a_p = a_p_u.*(am_p_u.<am_p_c) .+ a_p_c.*(am_p_u.>=am_p_c);
    mm.c = agrid.+mm.wages.-(1.0/Rf).*(mm.a_p.-mm.am_p).-(1.0/mm.Rm).*mm.am_p;

    mm.a_p = mm.a_p.*(mm.a_p.>0.0);
    mm.am_p = mm.am_p.*(mm.a_p.>0.0);
    mm.c = mm.c.*(mm.a_p.>0.0) .+ (agrid.+mm.wages).*(mm.a_p.<=0.0);
    for l=1:Nz
        for k=1:2
            for j=2:Na-1
                if mm.c[j,k,l]<mm.c[j-1,k,l]
                    # c[j,k]=(c[j-1,k]+c[j+1,k])/2.0
                    mm.c[j,k,l]=mm.c[j-1,k,l]+1e-10
                end
            end
            if mm.c[Na,k,l]<mm.c[Na-1,k,l]
                mm.c[Na,k,l]=mm.c[Na-1,k,l]+1e-10
            end
        end
    end

    #################################################
    #Update guess of Wa_p
    #################################################
    mm.Wa = u_prime.(mm.c;parms)
    return mm.Wa, mm.am_p, mm.a_p, mm.c, mm.ystar
end

###############################################################################
# Functions used in computing household decision rules and Wa
###############################################################################
#update ystar (unconstrained early-stage consumption)
function ystar_update(py,Wa_p,agrid;parms::parmsmodel)
    """
    Purpose: Given py_{t+1} and W'_{t+1}(a), update ystar_{t+1}
    inputs: py, Wa_p 
    output: ystar
    """
    ystar = zeros((Na,2,Nz));
    
    for k=1:2
        for j=1:Nz
            Wa_p_inv_int = LinearInterpolation(Wa_p[end:-1:1,k,j],Array(agrid)[end:-1:1],extrapolation_bc=Line());
            ygrid=range(1e-30,agrid[end]/py-1e-30,length=5000);
            atemp=Wa_p_inv_int.(υ_prime.(ygrid;parms)/py).+py*ygrid;
            ystar_int = LinearInterpolation(atemp,ygrid,extrapolation_bc=Line());
            ystar[:,k,j] = max.(min.(ystar_int.(agrid),(agrid.-1e-15)./py),zeros(Na))
        end
    end
    return ystar
end

#compute realized early-stage consumption given asset portfolios and liquidity shock
#only money meetings
function ym(ystar_fun,a,am,py)
    """
    Purpose: compute realized y given portfolio and price, only money accepted
    """
    return min.(ystar_fun.(a),am./py)
end
#all wealth meetings
function ymf(ystar_fun,a,am,py)
    """
    Purpose: compute realized y given portfolio and price, all wealth accepted
    """
    return min.(ystar_fun.(a),a./py).*transpose(ones(length(am)))
end

#update partials Vaj
function Vaj_update(py,ystar,Wa_p,agrid;parms::parmsmodel)
    """
    Purpose: compute partials Vaj_p, needed to compute Wa
    inputs: py, ystar, Wa_p 
    output: Vam_p, Vaf_p
    """
    #for over employment status
    Vam_p = zeros((Na,Na,2,Nz)); Vaf_p = zeros((Na,Na,2,Nz));

    for k=1:Nz
        for j=1:2
            #interpolate Wa_p
            Wa_p_int = LinearInterpolation(agrid, Wa_p[:,j,k], extrapolation_bc=Line());
            ystar_int = LinearInterpolation(agrid, ystar[:,j,k], extrapolation_bc=Line());
            temp = parms.ALPHA*(parms.ALPHAmf*(υ_prime.(ymf(ystar_int,agrid,transpose(agrid),py);parms)/py)).+(1.0-parms.ALPHA)*Wa_p[:,j,k].*transpose(ones(Na))
            Vam_p[:,:,j,k] = parms.ALPHA*((1.0-parms.ALPHAmf)*(υ_prime.(ym(ystar_int,agrid,transpose(agrid),py);parms)/py))+temp;
            Vaf_p[:,:,j,k] = parms.ALPHA*((1.0.-parms.ALPHAmf)*Wa_p_int.(agrid.-py*ym(ystar_int,agrid,transpose(agrid),py)))+temp;
        end
    end
    return Vam_p, Vaf_p
end

#update am_p_temp and c_temp
function policy_update_step1(Xam,Xaf,agrid;parms::parmsmodel)
    """
    Purpose: Given expected future utility first solve for policies (am',c) as functions of a'
    inputs: Xam, Xaf 
    outputs: am_p_temp, c_temp
    """
    #solve for am_prime and c as a function of a_prime
    am_p_temp = zeros((Na,2,Nz)); c_temp = zeros((Na,2,Nz));
    for l=1:Nz
        for k=1:2
            for j=1:Na
                Xam_int = LinearInterpolation(agrid,Xam[j,:,k,l],extrapolation_bc=Line());
                Xaf_int = LinearInterpolation(agrid,Xaf[j,:,k,l],extrapolation_bc=Line());
                ftemp(x) = Xam_int(x).-Xaf_int(x);
                if ftemp(agrid[1])*ftemp(agrid[j])<0
                    am_p_temp[j,k,l] = find_zero(ftemp,(agrid[1],agrid[j]),Roots.A42());
                elseif ftemp(agrid[j])>0
                    am_p_temp[j,k,l] = agrid[j];
                else
                    am_p_temp[j,k,l] = agrid[1];
                end
                c_temp[j,k,l] = u_prime_inv(Xam_int(am_p_temp[j,k,l]);parms);
            end
        end
    end
    return am_p_temp, c_temp
end

#update af_p, am_p, c 
function policy_update_step2(Rf,am_p_temp,c_temp,agrid; parms::parmsmodel,mm::model_outcomes)
    """
    purpose: update all policies 
    inputs: am_p_temp, c_temp, Rf, Rm
    outputs: policies am_p, a_p, c
    """
    am_p = zeros((Na,2,Nz)); a_p = zeros((Na,2,Nz)); c = zeros((Na,2,Nz));
    for j=1:2
        for k=1:Nz
            afun = c_temp[:,j,k].-mm.wages_bar[j,k].-mm.tau[:,j,k].+(1.0/Rf).*(agrid.-am_p_temp[:,j,k]).+(1.0/mm.Rm).*am_p_temp[:,j,k];
            a_p_fun = LinearInterpolation(afun, agrid, extrapolation_bc=Line());
            a_p[:,j,k] = max.(a_p_fun.(agrid),agrid[1]);
            am_p_fun = LinearInterpolation(agrid, am_p_temp[:,j,k], extrapolation_bc=Line());
            am_p[:,j,k] = am_p_fun.(a_p[:,j,k]);
            c_fun = LinearInterpolation(agrid, c_temp[:,j,k], extrapolation_bc=Line());
            c[:,j,k] = c_fun(a_p[:,j,k]);
        end
    end
    return am_p, a_p, c
end

#update a_p_temp and c_temp for constrained case am_p=a_p
function policy_update_constrained(Xam,agrid;parms::parmsmodel)
    """
    Purpose: update policies in constrained case of no financial wealth am_p=a_p
    """
    #solve for am_prime as a function of af_prime
    a_p = zeros((Na,2,Nz)); c = zeros((Na,2,Nz)); agrid_p=zeros((Na,2,Nz))
    for j=1:2
        for k=1:Nz
            c[:,j,k] = min.(u_prime_inv.(diag(Xam[:,:,j,k]);parms), agrid.+mm.wages_bar[j,k].+mm.tau[:,j,k].-(1.0/mm.Rm).*agrid[1])
            agrid_p[:,j,k] = c[:,j,k].+(1.0/mm.Rm).*agrid.-mm.wages_bar[j,k].-mm.tau[:,j,k]
            a_p_fun = LinearInterpolation(agrid_p[:,j,k],agrid,extrapolation_bc=Line());
            a_p[:,j,k] = a_p_fun.(agrid)
        end
    end
    am_p = copy(a_p)
    return am_p, a_p, c
end

###############################################################################
# Compute transition matrix given decision rules
###############################################################################
function transition_g!(mm::model_outcomes;py,parms::parmsmodel)
    """
    Purpose: update transition matrix T (Na*2,1) matrix on stacked mm.g = [mm.g0, mm.g1]
    inputs: policies ystar_{t+1}, a'_t, am'_t, c_t, PI_{t+1}
    output: T_t -> from distribution g_t to distribution g_{t+1}
    """
    mm.T = zeros((Na*2*Nz,Na*2*Nz));
    for l=0:Nz-1 #by productivity
        for k=0:1 #last stage employment state
            for kk=0:1 #beginning of next period emp state
                ystar_int = LinearInterpolation(mm.agrid, mm.ystar[:,kk+1,l+1], extrapolation_bc=Line());
                for j=1:Na #asset position
                    #no liquidity shock
                    ind = searchsortedlast(mm.agrid,mm.a_p[j,k+1,l+1]);
                    if ind<=0
                        ρ = 0.0;
                        mm.T[1 + kk*Na + Na*2*l,j+k*Na+ Na*2*l] += (1.0-ρ)*mm.PI'[kk+1,k+1]*(1.0-parms.ALPHA);
                    elseif ind>=Na
                        ρ = 1.0;
                        mm.T[Na + kk*Na + Na*2*l,j+k*Na + Na*2*l] += ρ*mm.PI'[kk+1,k+1]*(1.0-parms.ALPHA);
                    else
                        ρ = (mm.a_p[j,k+1,l+1]-mm.agrid[ind])/(mm.agrid[ind+1]-mm.agrid[ind]);
                        mm.T[ind+Na*kk + Na*2*l,j+k*Na + Na*2*l] += (1.0-ρ)*mm.PI'[kk+1,k+1]*(1.0-parms.ALPHA);
                        mm.T[(ind+1)+Na*kk + Na*2*l,j+k*Na + Na*2*l] += ρ*mm.PI'[kk+1,k+1]*(1.0-parms.ALPHA);
                    end
                    
                    #liquidity shock, only money
                    ind = searchsortedlast(mm.agrid,mm.a_p[j,k+1,l+1]-py*ym(ystar_int,mm.a_p[j,k+1,l+1],mm.am_p[j,k+1,l+1],py)[1])
                    if ind<=0
                        ρ = 0.0;
                        mm.T[1+kk*Na + Na*2*l,j+k*Na + Na*2*l] += (1.0-ρ)*mm.PI'[kk+1,k+1]*parms.ALPHA*(1.0-parms.ALPHAmf);
                    elseif ind>=Na
                        ρ = 1.0
                        mm.T[Na+kk*Na + Na*2*l,j+k*Na + Na*2*l] += ρ*mm.PI'[kk+1,k+1]*parms.ALPHA*(1.0-parms.ALPHAmf);
                    else
                        ρ = (mm.a_p[j,k+1,l+1]-py*ym(ystar_int,mm.a_p[j,k+1,l+1],mm.am_p[j,k+1,l+1],py)[1]-mm.agrid[ind])/(mm.agrid[ind+1]-mm.agrid[ind]);
                        mm.T[ind+Na*kk + Na*2*l,j+k*Na + Na*2*l] += (1.0-ρ)*mm.PI'[kk+1,k+1]*parms.ALPHA*(1.0-parms.ALPHAmf);
                        mm.T[(ind+1)+Na*kk + Na*2*l,j+k*Na + Na*2*l] += ρ*mm.PI'[kk+1,k+1]*parms.ALPHA*(1.0-parms.ALPHAmf);
                    end
                    
                    #liquidity shock, money and financial assets
                    ind = searchsortedlast(mm.agrid,mm.a_p[j,k+1,l+1]-py*ymf(ystar_int,mm.a_p[j,k+1,l+1],mm.am_p[j,k+1,l+1],py)[1])
                    if ind<=0
                        ρ = 0.0
                        mm.T[1+kk*Na + Na*2*l,j+k*Na + Na*2*l] += (1.0-ρ)*mm.PI'[kk+1,k+1]*parms.ALPHA*parms.ALPHAmf;
                    elseif ind>=Na
                        ρ = 1.0
                        mm.T[Na+kk*Na + Na*2*l,j+k*Na + Na*2*l] += ρ*mm.PI'[kk+1,k+1]*parms.ALPHA*parms.ALPHAmf;
                    else
                        ρ = (mm.a_p[j,k+1,l+1]-py*ymf(ystar_int,mm.a_p[j,k+1,l+1],mm.am_p[j,k+1,l+1],py)[1]-mm.agrid[ind])/(mm.agrid[ind+1]-mm.agrid[ind]);
                        mm.T[ind+Na*kk + Na*2*l,j+k*Na + Na*2*l] += (1.0-ρ)*mm.PI'[kk+1,k+1]*parms.ALPHA*parms.ALPHAmf;
                        mm.T[(ind+1)+Na*kk + Na*2*l,j+k*Na + Na*2*l] += ρ*mm.PI'[kk+1,k+1]*parms.ALPHA*parms.ALPHAmf;
                    end
                end
            end
        end
    end
end

###############################################################################
# Computes additional moments
###############################################################################
function other_mmts(; mm::model_outcomes, parms::parmsmodel)
    ##################################################
    #compute total wealth distribution    
    ##################################################
    G=zeros(Na); G0=zeros(Na); G1=zeros(Na);
    G[1]=sum(mm.g[1,:,:]);
    G1[1]=sum(mm.g[1,1,:]);
    G0[1]=sum(mm.g[1,2,:]);
    for i=2:Na
        G[i] = G[i-1] + sum(mm.g[i,:,:]);
        G1[i] = G1[i-1] + sum(mm.g[i,1,:]);
        G0[i] = G0[i-1] + sum(mm.g[i,2,:]);
    end

    ##################################################
    #liquid wealth to income distribution
    ##################################################
    Ntemp=100;
    gmi = zeros(Ntemp);
    li_range = range(1e-10,maximum(mm.am_p./(12.0.*mm.wages)),length=Ntemp);
    for k=1:Nz
        for j=1:2
            for i=1:Na
                temp=mm.am_p[i,j,k]/(12.0*mm.wages[i,j,k]);
                ind = searchsortedlast(li_range,temp);
                if ind<=0
                    ρ=0.0;
                    gmi[1]+=mm.g[i,j,k]
                elseif ind>=Ntemp
                    gmi[Ntemp]+=mm.g[i,j,k]
                else
                    ρ = (temp-li_range[ind])/(li_range[ind+1]-li_range[ind])
                    gmi[ind]+=(1-ρ)*mm.g[i,j,k]
                    gmi[ind+1]+=ρ*mm.g[i,j,k]
                end
            end
        end
    end
    Gmi=zeros(Ntemp);
    Gmi[1]=gmi[1];
    for i=2:Ntemp
        Gmi[i]=Gmi[i-1]+gmi[i]
    end
    
    ##################################################
    #liquid share distribution
    ##################################################
    Ntemp=100;
    gls = zeros(Ntemp);
    ls_range = range(1e-10,maximum(mm.am_p./mm.a_p),length=Ntemp);
    for k=1:Nz
        for j=1:2
            for i=1:Na
                temp=mm.am_p[i,j,k]/mm.a_p[i,j,k];
                ind = searchsortedlast(ls_range,temp);
                if ind<=0
                    ρ=0.0;
                    gls[1]+=mm.g[i,j,k]
                elseif ind>=Ntemp
                    gls[Ntemp]+=mm.g[i,j,k]
                else
                    ρ = (temp-ls_range[ind])/(ls_range[ind+1]-ls_range[ind])
                    gls[ind]+=(1-ρ)*mm.g[i,j,k]
                    gls[ind+1]+=ρ*mm.g[i,j,k]
                end
            end
        end
    end
    Gls = zeros(Ntemp);
    Gls[1]=gls[1];
    for i=2:Ntemp
        Gls[i]=Gls[i-1]+gls[i]
    end
    ind = searchsortedlast(ls_range,minimum(ls_range[Gls.>0])); #find nonzero ls
    ls_range2 = ls_range.*(1-ls_range[ind+1]) .+ ls_range[ind+1]
    ls_dist=LinearInterpolation(ls_range,Gls,extrapolation_bc=Line());
    Gls = ls_dist.(ls_range2)

    ##################################################
    #total wealth to income distribution 
    ##################################################
    Ntemp=100;
    gwi = zeros(Ntemp);
    wi_range = range(1e-10,maximum(mm.a_p./(12.0.*mm.wages)),length=Ntemp);
    for k=1:Nz
        for j=1:2
            for i=1:Na
                temp=mm.a_p[i,j,k]/(12.0*mm.wages[i,j,k]);
                ind = searchsortedlast(wi_range,temp);
                if ind<=0
                    ρ=0.0;
                    gwi[1]+=mm.g[i,j,k]
                elseif ind>=Ntemp
                    gwi[Ntemp]+=mm.g[i,j,k]
                else
                    ρ = (temp-wi_range[ind])/(wi_range[ind+1]-wi_range[ind])
                    gwi[ind]+=(1-ρ)*mm.g[i,j,k]
                    gwi[ind+1]+=ρ*mm.g[i,j,k]
                end
            end
        end
    end
    Gwi=zeros(Ntemp);
    Gwi[1]=gwi[1];
    for i=2:Ntemp
        Gwi[i]=Gwi[i-1]+gwi[i]
    end

    ##################################################
    #Gini coefficient wrt total wealth
    ##################################################
    ghat=sum(mm.g,dims=(2,3))[:,1,1];
    gini=0.0
    for i=1:Na 
        for j=1:Na 
            gini+= ghat[i]*ghat[j]*abs(mm.agrid[i]-mm.agrid[j])
        end
    end
    gini=gini/(2*sum(ghat.*mm.agrid))
    
    ##################################################
    #Other aggregate moments
    ##################################################
    liquid_share = mm.Zmd/(mm.Jd+mm.Zmd);
    liquid_to_income = mm.Zmd/(12*sum(mm.g.*mm.wages));
    gbonds_share = mm.Ag/(mm.Jd+mm.Zmd);
    equity_share = (mm.emp*sum(mm.gz.*mm.Js))/(mm.Jd+mm.Zmd);
    wealth_to_income = (mm.Jd+mm.Zmd)/(12*sum(mm.g.*mm.wages));
    liquid_prem = ((mm.Rf^12)-1.0)*100 - ((mm.Rm^12)-1.0)*100;

    return G, G1, G0, gmi, Gmi, li_range, gls, Gls, ls_range, gwi, Gwi, wi_range, gini, liquid_share, liquid_to_income, gbonds_share, equity_share, wealth_to_income, liquid_prem
end

#END